#!/usr/bin/env python3

"""
author: liuxu
date: 2025/9/3
description: analyze_trace.sh的python版本，
    统计hip-trace的空泡时间，支持跨文件合并相同TID数据
"""

import json
import sys
import os
import re
import glob
from collections import defaultdict


def main():
    if len(sys.argv) < 2:  # file_pattern： result_*.json
        print(f"Usage: {sys.argv[0]} <file_pattern>", file=sys.stderr)
        sys.exit(1)

    # 获取json文件列表（支持通配符）
    json_files = []
    for pattern in sys.argv[1:]:
        json_files.extend(glob.glob(pattern))
    # 提取json后缀数字排序，确保后续跨文件匹配push，pop按顺序出现
    json_files = sorted(json_files, key=lambda x: int(re.search(r"_(\d+)\.json$", x).group(1)))

    if not json_files:
        print("[ERROR] No matching file was found", file=sys.stderr)
        sys.exit(1)

    # 全局数据结构，用于跨文件存储数据
    all_tid_data = defaultdict(list)
    push_events = defaultdict(list)  # 保存未匹配的push事件

    # 事件名称常量定义
    PUSH_EVENT = '__hipPushCallConfiguration'
    POP_EVENT = '__hipPopCallConfiguration'

    # 处理所有json文件
    for file in json_files:
        if not os.path.isfile(file):
            print(f"Warning: File '{file}' not found, skipping...", file=sys.stderr)
            continue

        try:
            with open(file, 'r') as f:
                data = json.load(f)
        except json.JSONDecodeError as e:
            print(f"[ERROR] Invalid JSON in {file}: {str(e)}", file=sys.stderr)
            continue
        except Exception as e:
            print(f"[ERROR] Failed to parse {file}: {str(e)}", file=sys.stderr)
            continue

        # 处理当前文件的trace events（按事件在文件中的顺序处理）
        # events = sorted(data.get('traceEvents', []), key=lambda x: x.get('args', {}).get('index', -1))
        events = data.get('traceEvents', [])
        for event in events:
            event_name = event.get('name')
            if event_name not in (PUSH_EVENT, POP_EVENT):
                continue

            args = event.get('args', {})
            tid = args.get('tid')
            if tid is None:
                print(f"ERROR: Event {event_name} missing TID: {event}", file=sys.stderr)
                continue

            # 统一处理时间字段的大小写问题,25.10将字段首字母改为小写
            def get_time(key):
                return args.get(key) or args.get(key[0].lower() + key[1:])

            # if begin_ns is None or end_ns is None:
            #     print(f"ERROR: Abnormal event data (missing time) in {file}: {event}", file=sys.stderr)
            #     continue

            if event_name == PUSH_EVENT:
                push_end_ns = int(get_time('EndNs'))
                # 保存push事件的完整信息
                push_events[tid].append({'end_ns': push_end_ns})
                # push_events[tid].append({
                #     'name': event_name,
                #     'begin_ns': begin_ns,
                #     'end_ns': end_ns,
                #     'file': file,
                #     'args': args
                # })
            elif event_name == POP_EVENT:
                pop_begin_ns = int(get_time('BeginNs'))
                # 按顺序匹配最近的push事件
                if not push_events[tid]:
                    print(f"ERROR: Orphan Pop event in TID {tid}", file=sys.stderr)
                    print(f"  Pop event begins at {pop_begin_ns}ns in {file}", file=sys.stderr)
                    continue

                # 取出上一个未匹配的Push事件（按顺序配对）
                push_event = push_events[tid].pop(-1)
                # 计算差值：Pop的begin_ns - Push的end_ns
                time_diff = pop_begin_ns - push_event['end_ns']

                if time_diff >= 0:
                    all_tid_data[tid].append(time_diff)
                else:
                    print(f"ERROR: Negative time difference for TID {tid}: {time_diff}ns", file=sys.stderr)
                    print(f"  Push event ends at {push_event['end_ns']}ns in {push_event['file']}", file=sys.stderr)
                    print(f"  Pop event begins at {pop_begin_ns}ns in {file}", file=sys.stderr)

    # 处理所有文件后，检查是否有未匹配的push事件
    orphan_count = 0
    for tid, events in push_events.items():
        if events:
            orphan_count += len(events)
            for push_event in events:
                print(f"ERROR: Orphan Push event in TID {tid}")
                print(f"  Event: {push_event['name']} at {push_event['end_ns']} in {push_event['file']}")
                # 打印参数信息
                args = push_event.get('args', {})
                if args:
                    print(f"  Args: {json.dumps(args, indent=2)}")

    # 计算所有TID的统计信息
    tid_stats = {}
    for tid, diffs in all_tid_data.items():
        if not diffs:
            continue

        count = len(diffs)
        total = sum(diffs)
        min_val = min(diffs)
        max_val = max(diffs)
        first_val = diffs[0]

        # 计算第二最大值
        second_max = max_val
        if count > 1:
            # 优化第二最大值计算，避免完整排序
            max1 = max2 = -float('inf')
            for val in diffs:
                if val > max1:
                    max2 = max1
                    max1 = val
                elif val > max2:
                    max2 = val
            second_max = max2

        # 计算平均值
        avg = total / count

        # 计算排除第一个数据后的平均值
        avg_excl_first = sum(diffs[1:]) / (count - 1) if count > 1 else 0

        tid_stats[tid] = {
            'count': count,
            'first': first_val,
            'max': max_val,
            'second_max': second_max,
            'min': min_val,
            'avg': avg,
            'avg_excl_first': avg_excl_first
        }

    # 打印结果 - 保持左对齐固定列宽格式
    print(f"{'='*10} Analyzed bubble time of files: ({', '.join(sys.argv[1:])}) {'='*10}")

    header_format = (
        "{:<8}  {:<10}  {:<12}  {:<14}  {:<14}  {:<12}  {:<12}  {:<18}"
    )
    data_format = (
        "{:<8}  {:<10}  {:<12}  {:<14}  {:<14}  {:<12}  {:<12.2f}  {:<18.2f}"
    )

    # 打印表头
    print(header_format.format(
        "TID", "Count", "First(ns)", "Max(ns)",
        "SecondMax(ns)", "Min(ns)", "Avg(ns)", "AvgExclFirst(ns)"
    ))

    # 按TID排序输出数据
    for tid in sorted(tid_stats.keys(), key=lambda x: int(x) if str(x).isdigit() else x):
        stats = tid_stats[tid]
        print(data_format.format(
            tid,
            stats['count'],
            stats['first'],
            stats['max'],
            stats['second_max'],
            stats['min'],
            stats['avg'],
            stats['avg_excl_first']
        ))

    if orphan_count > 0:
        print(f"ERROR: Total orphan events: {orphan_count} push events", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
